from ultralytics import YOLO
import cv2 # 用于摄像头或视频流的额外处理，如果需要

# --- 1. 模型加载与基本配置 ---
# 模型的路径。可以是本地 .pt 文件，也可以是 Ultralytics 预训练模型的名称 (如 'yolov8n.pt')
# 如果是分类、分割、姿态模型，请加载对应的模型文件，例如 'yolov8n-seg.pt', 'yolov8n-pose.pt'
MODEL_PATH = '/home/youngqi/Public/ultralytics_yolov8/runs/detect/YOLOv8_test/train/weights/best.pt'  

# 加载 YOLO 模型
# 根据你的任务（检测、分割、姿态、分类）选择合适的模型文件
# model = YOLO(MODEL_PATH, task='detect') # 显式指定任务，通常不需要，YOLO会自动识别
model = YOLO(MODEL_PATH)  

print(f"模型 '{MODEL_PATH}' 已成功加载。")

# --- 2. 推理参数配置 (高度可自定义) ---
# 这些参数将传递给 model.predict() 方法，控制推理行为。
# 你可以根据需要调整这些值。
INFERENCE_PARAMS = {
    # 图像大小 (imgsz): 用于推理的图像尺寸。
    # 可以是单个整数（如 640，表示 640x640），也可以是元组（如 (1280, 720)）。
    # 较大的尺寸可以更好地检测小目标，但会增加计算量和内存消耗。
    # 你的数据集是 1920x1080，可以尝试 1280 或更高的值，但要确保 GPU 内存足够。
    'imgsz': 640, # 默认值，可以尝试 960, 1280, (1920, 1080) 等

    # 置信度阈值 (conf): 只有当预测的置信度高于此值时，才显示该预测。
    # 范围 0.0 到 1.0。值越低，检测到的目标越多（可能包含更多假阳性）。
    'conf': 0.25, # 默认 0.25，可以调高到 0.5 或更低到 0.1 

    # IoU 阈值 (iou): 用于非极大值抑制 (NMS) 的交并比阈值。
    # 较高的值会使更多的重叠框被保留。范围 0.0 到 1.0。
    'iou': 0.7, # 默认 0.7，可以调低到 0.5 或调高到 0.9 

    # 设备 (device): 指定运行推理的设备。
    # 'cpu' 表示使用 CPU，0 表示使用第一个 GPU，0,1,2,3 表示使用多个 GPU。
    # 如果你有 GPU，强烈建议设置为 0。
    'device': '0', # 'cpu' 或 '0' 

    # 显示结果 (show): 是否在屏幕上实时显示推理结果。
    'show': True, # True 或 False 

    # 保存结果 (save): 是否将结果保存到 runs/detect/predictX 文件夹。
    # 结果包括带有边界框的图像/视频。
    'save': False, # True 或 False 

    # 保存检测文本 (save_txt): 是否将边界框坐标保存为 TXT 文件。
    # 对于每个检测到的对象，会保存其类别、归一化坐标。
    'save_txt': False, # True 或 False 

    # 保存置信度到 TXT (save_conf): 配合 save_txt 使用，在 TXT 文件中包含置信度分数。
    'save_conf': False, # True 或 False 

    # 保存裁剪图像 (save_crop): 是否将检测到的对象从原始图像中裁剪出来并保存。
    'save_crop': False, # True 或 False 

    # 显示标签 (show_labels): 在结果图像上显示类别标签。
    'show_labels': True, # True 或 False 

    # 显示置信度 (show_conf): 在结果图像上显示置信度分数。
    'show_conf': True, # True 或 False 

    # 显示边界框 (show_boxes): 在结果图像上显示边界框。
    'show_boxes': True, # True 或 False 
     
    # 边界框线条宽度 (line_width): 设置绘制边界框的线条粗细。
    # None 表示默认值，可以是整数如 2, 3 等。
    'line_width': None,  

    # 类别过滤 (classes): 过滤要检测的特定类别。
    # 可以是单个整数（如 0 表示第一个类别），或整数列表（如 [0, 2, 5]）。
    # None 表示检测所有类别。
    # 假设你的 iflytek_ucar_phone.yaml 定义了 5 个类别，你可以根据索引来选择。
    # 例如，如果类别 "phone" 的索引是 0，你可以设置为 0。
    'classes': None, # 例如: 0 (只检测第一个类别), [0, 1] (检测前两个类别) 

    # 类别无关的 NMS (agnostic_nms): 是否执行类别无关的非极大值抑制。
    # True 表示不同类别的重叠框也会被抑制，False 则只在同一类别内抑制。
    'agnostic_nms': False, # True 或 False 

    # 跟踪器 (tracker): 启用对象跟踪功能。
    # 需要在模型加载时就支持跟踪，或者是一个跟踪器配置文件路径。
    # 例如 'bytetrack.yaml' 或 'botsort.yaml'。
    # 如果不进行跟踪，设置为 None。
    'tracker': None, # 例如 'bytetrack.yaml', 'botsort.yaml', 或 None 

    # Retina Mask (retina_masks): 仅用于分割模型，生成高分辨率掩码。
    'retina_masks': False, # True 或 False，仅当任务为 'segment' 时有效 
} 

# --- 3. 输入源选择与测试执行 --- 

# 定义你的输入源。你可以注释掉不需要的，只保留一个或多个进行测试。
# 每次只取消注释一个 TEST_ITEM 进行测试，或者按顺序执行。

# --- 测试项 1: 视频文件推理 ---
# 注意: 如果视频较长，show=True 可能导致性能问题，可以考虑设置为 False 并 save=True。
# TEST_ITEM_VIDEO = {
#     'name': "视频文件测试",
#     'source': '/home/youngqi/Videos/camera_video/robot_video_20250625_114645.mp4', # 替换为你的视频文件路径，例如 'data/videos/test.mp4'
#     'params': INFERENCE_PARAMS.copy() # 复制一份通用参数，方便针对此项进行修改
# }
# # 可以针对此测试项修改特定参数，例如增大 imgsz 或保存结果
# TEST_ITEM_VIDEO['params']['imgsz'] = 640
# TEST_ITEM_VIDEO['params']['save'] = True
# TEST_ITEM_VIDEO['params']['show'] = True # 保持显示
# print(f"\n--- 正在执行: {TEST_ITEM_VIDEO['name']} ---")
# # 如果 source 是视频文件，results 将是一个生成器，你可以遍历它
# results = model.predict(**TEST_ITEM_VIDEO['params'], source=TEST_ITEM_VIDEO['source'])
# for r in results:
#     # 如果你需要对每一帧结果进行额外处理，可以在这里添加代码
#     pass
# print(f"{TEST_ITEM_VIDEO['name']} 完成。")


# --- 测试项 2: 摄像头实时推理 ---
# 请确保你的摄像头设备 ID 正确 (通常 0 是默认内置摄像头，1 是外接摄像头等)。
# 'show' 参数必须为 True 才能看到摄像头画面。
# 注意: 实时推理对性能要求较高，根据你的 GPU 调整 imgsz 和 batch 大小。

# TEST_ITEM_WEBCAM = {
#     'name': "摄像头实时测试",
#     'source': 0, # 摄像头 ID，0 为默认摄像头
#     'params': INFERENCE_PARAMS.copy()
# }
# TEST_ITEM_WEBCAM['params']['imgsz'] = 640 # 摄像头通常用较小尺寸以保证流畅性
# TEST_ITEM_WEBCAM['params']['show'] = True # 必须显示才能看到实时画面
# TEST_ITEM_WEBCAM['params']['save'] = False # 通常不保存摄像头实时流
# print(f"\n--- 正在执行: {TEST_ITEM_WEBCAM['name']} ---")
# # 摄像头推理也是一个生成器
# results = model.predict(**TEST_ITEM_WEBCAM['params'], source=TEST_ITEM_WEBCAM['source'])
# for r in results:
#     # 如果你需要对每一帧结果进行额外处理，可以在这里添加代码
#     # 例如：你可以从 r.plot() 获取带有检测框的图像帧
#     # frame = r.plot()
#     # cv2.imshow('YOLOv8 Webcam', frame)
#     # if cv2.waitKey(1) & 0xFF == ord('q'): # 按 'q' 键退出
#     #     break
# # cv2.destroyAllWindows() # 关闭所有 OpenCV 窗口
#     print(f"{TEST_ITEM_WEBCAM['name']} 完成。")


# --- 测试项 3: 单张图片推理 ---
# 图片推理会更快完成，结果会保存到 runs/detect/predictX 文件夹（如果 save=True）。
# TEST_ITEM_IMAGE = {
#     'name': "单张图片测试",
#     'source': 'your_image.jpg', # 替换为你的图片文件路径，例如 'data/images/bus.jpg'
#     'params': INFERENCE_PARAMS.copy()
# }
# TEST_ITEM_IMAGE['params']['imgsz'] = 1280 # 图片可以尝试更高的分辨率
# TEST_ITEM_IMAGE['params']['save'] = True # 保存检测结果图片
# TEST_ITEM_IMAGE['params']['show'] = True # 显示结果图片
# print(f"\n--- 正在执行: {TEST_ITEM_IMAGE['name']} ---")
# # 对于单张图片，results 是一个列表，每个元素是 Result 对象
# results = model.predict(**TEST_ITEM_IMAGE['params'], source=TEST_ITEM_IMAGE['source'])
# # result = results[0] # 获取第一个 (也是唯一一个) 结果对象
# # print(f"检测到的类别和数量: {result.boxes.cls.unique()}") # 打印检测到的类别
# print(f"{TEST_ITEM_IMAGE['name']} 完成。")


# --- 测试项 4: 图片文件夹推理 ---
# 会批量处理文件夹内的所有图片。
TEST_ITEM_IMAGE_FOLDER = {
    'name': "图片文件夹测试",
    'source': 'datasets/iflytek_phone_ucar/', # 替换为你的图片文件夹路径，例如 'data/images/'
    'params': INFERENCE_PARAMS.copy()
}
TEST_ITEM_IMAGE_FOLDER['params']['imgsz'] = 640  
TEST_ITEM_IMAGE_FOLDER['params']['save'] = True  
TEST_ITEM_IMAGE_FOLDER['params']['show'] = False # 批量处理时通常不显示，否则会弹出大量窗口
print(f"\n--- 正在执行: {TEST_ITEM_IMAGE_FOLDER['name']} ---")
# 文件夹推理也是一个生成器
results = model.predict(**TEST_ITEM_IMAGE_FOLDER['params'], source=TEST_ITEM_IMAGE_FOLDER['source'])
for r in results:
    # 每处理一张图片，r 都是一个 Result 对象
    pass
print(f"{TEST_ITEM_IMAGE_FOLDER['name']} 完成。")


# --- 测试项 5: YouTube 视频流推理 (需要网络连接) ---
# 注意: 这需要你的网络连接稳定，且可能会有延迟。
# TEST_ITEM_YOUTUBE = {
#     'name': "YouTube 视频流测试",
#     'source': 'https://www.youtube.com/watch?v=your_video_id', # 替换为有效的 YouTube 视频 ID 或完整 URL
#     'params': INFERENCE_PARAMS.copy()
# }
# TEST_ITEM_YOUTUBE['params']['imgsz'] = 640  
# TEST_ITEM_YOUTUBE['params']['show'] = True  
# TEST_ITEM_YOUTUBE['params']['save'] = False  
# print(f"\n--- 正在执行: {TEST_ITEM_YOUTUBE['name']} ---")
# results = model.predict(**TEST_ITEM_YOUTUBE['params'], source=TEST_ITEM_YOUTUBE['source'])
# for r in results:
#     pass
# print(f"{TEST_ITEM_YOUTUBE['name']} 完成。")


# --- 测试项 6: RTSP 视频流推理 (需要网络摄像头或NVR) ---
# 通常用于工业监控或网络摄像头。
# TEST_ITEM_RTSP = {
#     'name': "RTSP 视频流测试",
#     'source': 'rtsp://user:password@ip_address:port/stream', # 替换为你的 RTSP 地址
#     'params': INFERENCE_PARAMS.copy()
# }
# TEST_ITEM_RTSP['params']['imgsz'] = 640  
# TEST_ITEM_RTSP['params']['show'] = True  
# TEST_ITEM_RTSP['params']['save'] = False  
# print(f"\n--- 正在执行: {TEST_ITEM_RTSP['name']} ---")
# results = model.predict(**TEST_ITEM_RTSP['params'], source=TEST_ITEM_RTSP['source'])
# for r in results:
#     pass
# print(f"{TEST_ITEM_RTSP['name']} 完成。")


# print("\n所有选定的测试项已执行完毕。")
# print("结果通常保存在 'runs/detect/' 或 'runs/segment/' 等目录下（取决于任务和保存设置）。")